import argparse
import socket
import struct

def send_ip_packet(dst_ip, message):
    # 创建一个 UDP 套接字
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    # 发送消息
    sock.sendto(message.encode(), (dst_ip, 12345))
    sock.close()

def send_source_routing_packet(dst_ip, message, ports):
    # 构建源路由头
    source_route_header = b''
    for port in ports:
        bos = 1 if port == ports[-1] else 0
        source_route_header += struct.pack('!BB', bos, port)

    # 创建一个 UDP 套接字
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    # 发送包含源路由头的消息
    packet = source_route_header + message.encode()
    sock.sendto(packet, (dst_ip, 12345))
    sock.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Send IP or Source Routing Packet')
    parser.add_argument('dst_ip', type=str, help='Destination IP address')
    parser.add_argument('message', type=str, help='Message to send')
    parser.add_argument('--source_routing', action='store_true', help='Send source routing packet')
    parser.add_argument('--ports', nargs='+', type=int, help='List of ports for source routing')

    args = parser.parse_args()

    if args.source_routing:
        if not args.ports:
            print("Please provide a list of ports for source routing.")
        else:
            send_source_routing_packet(args.dst_ip, args.message, args.ports)
    else:
        send_ip_packet(args.dst_ip, args.message)